#include "tensor.h"
#include <math.h> 
#include <map>

tensor::tensor() {}

tensor::tensor(double v) {
  shape_.push_back(1);
  data_.push_back(v);
}

//累乘求得数据个数
size_t tensor::getTotalSize() const {
  return std::accumulate(shape_.begin(), shape_.end(), (size_t)1,
                         std::multiplies<size_t>());
}

tensor::tensor(int dim, size_t shape[], double data[]) {
  std::copy(shape, shape + dim, std::back_inserter(shape_));
  std::copy(data, data + getTotalSize(), std::back_inserter(data_));
}

int tensor::get_dim() const { return shape_.size(); }

size_t *tensor::get_shape_array() {
  return shape_.empty() ? nullptr : &shape_[0];
}

double *tensor::get_data_array() { return &data_[0]; }
std::vector<double>& tensor::get_data_vec(){
   return  this->data_;
};


template <class T> 
tensor helper(const tensor &lhs, const tensor &rhs) {
  tensor result;
  result.shape_ = lhs.shape_;
  size_t totalSize = lhs.getTotalSize();
  result.data_.reserve(totalSize);
  std::transform(lhs.data_.begin(), lhs.data_.begin() + totalSize,
                 rhs.data_.begin(), std::back_inserter(result.data_), T());
  return result;
}

tensor operator+(const tensor &lhs, const tensor &rhs) {
  return helper<std::plus<double>>(lhs, rhs);
}

tensor operator-(const tensor &lhs, const tensor &rhs) {
  return helper<std::minus<double>>(lhs, rhs);
}

tensor operator*(const tensor &lhs, const tensor &rhs) {
  tensor result;
  size_t lhsSize = lhs.getTotalSize();
  size_t rhsSize = rhs.getTotalSize();
  if (lhsSize == 1 || rhsSize == 1) {
    double scalar = (lhsSize == 1 ? lhs : rhs).data_[0];
    const tensor &t = lhsSize == 1 ? rhs : lhs;
    result.shape_ = t.shape_;
    result.data_.reserve(lhsSize * rhsSize);
    std::transform(t.data_.begin(), t.data_.end(),
                   std::back_inserter(result.data_),
                   [&scalar](double x) { return x * scalar; });
  } else {
    size_t I = lhs.shape_[0];
    size_t J = rhs.shape_[1];
    size_t K = lhs.shape_[1];
    result.shape_ = {I, J};
    size_t totalSize = result.getTotalSize();
    result.data_.resize(totalSize);

    // can be optimized
    for (size_t i = 0; i < I; ++i) {
      for (size_t j = 0; j < J; ++j) {
        for (size_t k = 0; k < K; ++k) {
          result.data_[J * i + j] +=
              lhs.data_[K * i + k] * rhs.data_[J * k + j];
        }
      }
    }
  }
  return result;
}

tensor ReLU(const tensor &x){
  tensor result;
  size_t xSize = x.getTotalSize();
  result.shape_ = x.shape_;
  for(size_t i = 0; i < x.getTotalSize(); i++ ){
      result.data_.push_back(x.data_[i]*(x.data_[i]>0));
  }
  return result; 
  }
  
tensor Flatten(const tensor &x){
  tensor result;
  size_t xSize = x.getTotalSize();
  size_t shape_1 = xSize/x.shape_[0];
  std::vector<size_t> shape_;
  shape_.push_back(x.shape_[0]);
  shape_.push_back(shape_1);
  result.shape_ = shape_;
  result.data_ = x.data_;
  return result;
  }
  
tensor Inut2d(const tensor &x){
     tensor result;
     size_t xSize = x.getTotalSize();
     std::vector<size_t> shape_;
     shape_.push_back(x.shape_[0]);
     shape_.push_back(x.shape_[3]);
     shape_.push_back(x.shape_[1]);
     shape_.push_back(x.shape_[2]);
     result.shape_ = shape_;
     for(size_t i = 0; i < x.shape_[0]; i++ ){
      for(size_t j =0; j < x.shape_[3];j++){
        for(size_t k = 0; k <x.shape_[1]*x.shape_[2]; k++){
            result.data_.push_back(x.data_[j + k*x.shape_[3] + i*x.shape_[3]*x.shape_[1]*x.shape_[2]]);
        }
      }        
    }
    return result;
}

tensor Linear(const tensor &w, const tensor &x, const tensor &b){
    //计算w^T
    tensor result_wT;
    std::vector<size_t> shape_0;
    shape_0.push_back(x.shape_[1]);
    shape_0.push_back(w.shape_[0]);
    size_t totalSize_w = w.getTotalSize();
    result_wT.data_.reserve(totalSize_w);
    result_wT.shape_ = shape_0;
    for(size_t i = 0; i < w.shape_[1]; i++){
      for(size_t j = 0; j < w.shape_[0]; j++){
         result_wT.data_.push_back(w.data_[i + j*w.shape_[1]]);
      }
    }
    tensor result_xwT;
    std::vector<size_t> shape_xwT;
    shape_xwT.push_back(x.shape_[0]);
    shape_xwT.push_back(result_wT.shape_[1]);
    size_t totalSize_xwT = x.shape_[0]*result_wT.shape_[1];
    result_xwT.data_.resize(totalSize_xwT);
    result_xwT.shape_ = shape_xwT;
    result_xwT =  x * result_wT; //计算 x*w^T

    tensor result_bBrodcast;    
    std::vector<size_t> shape_1;
    shape_1.push_back(b.shape_[0]);
    shape_1.push_back(1*x.shape_[0]);
    result_bBrodcast.shape_ = shape_1;
    for(size_t i = 0; i < result_bBrodcast.shape_[0]; i++){
       for(size_t j = 0; j < result_bBrodcast.shape_[1]; j++){
          result_bBrodcast.data_.push_back(b.data_[i]);
       }
    }

   //计算 b^T
   tensor result_bBroadcastT;
   std::vector<size_t> shape_2;
   size_t totalSize_b = result_bBrodcast.getTotalSize();
   result_bBrodcast.data_.resize(totalSize_b);
   shape_2.push_back(result_bBrodcast.shape_[1]);
   shape_2.push_back(result_bBrodcast.shape_[0]);
   result_bBroadcastT.shape_ = shape_2;
  for(size_t i = 0; i < result_bBrodcast.shape_[1]; i++){
      for(size_t j = 0; j < result_bBrodcast.shape_[0]; j++){
         result_bBroadcastT.data_.push_back(result_bBrodcast.data_[i + j*result_bBrodcast.shape_[1]]);
      }
    }
    tensor result; 
    result.data_.resize(result_xwT.shape_[0]*result_xwT.shape_[1]);
    result = result_xwT + result_bBroadcastT;  // 计算  x*w^T + b^T
    return result;
}





std::vector<double> maxpooling(std::vector<std::vector<double> > map, int k_w, int k_h, int s_w, int s_h)
{
	int row = map.size(); int col = map[0].size();
	int out_row = 1+(row-k_h)/s_h, out_col = 1+(col-k_w)/s_w;
	int mod_row = row%s_h, mod_col = col%s_w;
	if (mod_row != 0) out_row++;
	if (mod_col != 0) out_col++;

	std::vector<std::vector<double> > res(out_row, std::vector<double>(out_col, 0));
	std::vector<std::vector<double> > pad_map(map);

	for (int i=0;i<row;i++)
		for (int j = 0; j < k_h - mod_row; j++)
		{
			pad_map[i].push_back(map[i][col - 1]);
		}
	for (int j = 0; j < k_w - mod_col; j++)
	{
		pad_map.push_back(pad_map[row - 1]);
	}

	for(int i=0;i<out_row;i++)
		for (int j = 0; j < out_col; j++)
		{
			int start_x = j*s_w;
			int start_y = i*s_h;
			std::vector<double> temp;
			for(int ii=0;ii<k_w;ii++)
				for (int jj = 0; jj < k_h; jj++)
				{
					temp.push_back(pad_map[start_y + jj][start_x + ii]);
				}
			sort(temp.begin(), temp.end());
			res[i][j] = temp[temp.size() - 1];
		}
       std::vector<double> res_;
      //将二维数组转为一位数组     
      for (auto l : res)
 	     {
		       for (auto ll : l)
		       {
			          res_.push_back(ll);
		       }
	      }
	return res_;
}


tensor MaxPool2d(const tensor &x, const int kernel_size, const int stride){
    std::cout <<"======1================="<<std::endl;
    std::cout << x.shape_[0] <<"_"<< x.shape_[1] << x.shape_[2] << x.shape_[3] << std::endl;
    std::cout << kernel_size << std::endl;
    std::cout << stride << std::endl;
    std::cout <<"======2================="<<std::endl;

    tensor result;
    size_t xSize = x.getTotalSize();
    std::cout << xSize << std::endl;
    std::vector<size_t> shape_;
    shape_.push_back(x.shape_[0]);
    shape_.push_back(x.shape_[1]);
    shape_.push_back(floor(x.shape_[2]/kernel_size));
    shape_.push_back(floor(x.shape_[3]/kernel_size));
    result.shape_ = shape_;
    for(size_t i = 0; i < x.shape_[0]; i++ ){
      for(size_t j =0; j < x.shape_[1];j++){
        std::vector<std::vector<double> > tmp; 
        for(size_t k =0; k < x.shape_[2];k++){
          std::vector<double> tmp_;
          for(size_t l =0; l < x.shape_[3];l++){
            double int_data = x.data_[l + k*x.shape_[3] + j*x.shape_[3]*x.shape_[2] + i*x.shape_[3]*x.shape_[2]*x.shape_[1] ];
            tmp_.push_back(int_data);   
          }
          tmp.push_back(tmp_);
        }
       auto res = maxpooling(tmp,kernel_size,kernel_size,stride,stride);
       for(int i = 0; i< res.size(); i++){
         result.data_.push_back(res[i]);
       }
      }        
    }
    return result;   
  }












std::ostream &operator<<(std::ostream &os, const tensor &t) {
  std::ostream_iterator<double> it(os, " ");
  std::copy(t.shape_.begin(), t.shape_.end(), it);
  os << '\n';
  std::copy(t.data_.begin(), t.data_.end(), it);
  os << '\n';
  return os;
}








